package com.wispx.gateway.core.filter.loadblance.strategy;

import com.wispx.gateway.config.pojo.RouteDefinition;
import com.wispx.gateway.config.pojo.ServiceInstance;
import com.wispx.gateway.config.util.FilterConfigUtil;
import com.wispx.gateway.core.context.GatewayContext;

import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicInteger;

import static com.wispx.gateway.common.constant.FilterConstant.LOAD_BALANCE_FILTER_NAME;
import static com.wispx.gateway.common.constant.LoadBalanceConstant.ROUND_ROBIN_LOAD_BALANCE_STRATEGY;

/**
 * 轮询负载均衡
 */
public class RoundRobinLoadBalacneStrategy implements LoadBalanceStrategy {

    private Map<String/* 服务名 */, AtomicInteger> strictPositionMap = new ConcurrentHashMap<>();
    private Map<String/* 服务名 */, Integer> positionMap = new ConcurrentHashMap<>();
    private final int THRESHOLD = Integer.MAX_VALUE >> 2; // 访问次数的阈值，超过阈值重置为0

    @Override
    public ServiceInstance selectInstance(GatewayContext context, List<ServiceInstance> instances) {
        RouteDefinition.LoadBalanceFilterConfig loadBalanceFilterConfig = FilterConfigUtil.findFilterConfigByClass(context.getRouteDefinition().getFilterConfigs(), LOAD_BALANCE_FILTER_NAME, RouteDefinition.LoadBalanceFilterConfig.class);
        boolean isStrict = loadBalanceFilterConfig != null ? loadBalanceFilterConfig.isStrictRoundRobin() : true;
        String serviceName = context.getRequest().getServiceDefinition().getServiceName();
        if (isStrict) {
            AtomicInteger visits = strictPositionMap.computeIfAbsent(serviceName, name -> new AtomicInteger(0));
            ServiceInstance selectedInstance = instances.get(visits.get() % instances.size());
            int count = visits.incrementAndGet();
            if (count >= THRESHOLD) {
                visits.set((count + 1) % instances.size());
            }
            return selectedInstance;
        } else {
            int visits = positionMap.computeIfAbsent(serviceName, name -> 0);
            ServiceInstance selectedInstance = instances.get(visits % instances.size());
            positionMap.put(serviceName, ++visits);
            if (visits >= THRESHOLD) {
                positionMap.put(serviceName, visits % instances.size());
            }
            return selectedInstance;
        }
    }

    @Override
    public String mark() {
        return ROUND_ROBIN_LOAD_BALANCE_STRATEGY;
    }
}
